import pygame, typing
from uuid import UUID, uuid4 as new_uuid
from dataclasses import dataclass

@dataclass
class PointData:
    name: str
    type: typing.Literal['input', 'output']
    order: int

class Node:
    def __init__(self, x: int, y: int):
        self.x = x
        self.y = y

        self.points: set[PointData] = set()

        self.name = self.__class__.__name__

    def draw(self, surface: pygame.Surface, color: pygame.Color, highlight_color: pygame.Color, text_color: pygame.Color, highlight=False, x_offset: int = 0, y_offset: int = 0):
        pygame.draw.rect(surface, color, (self.x + x_offset, self.y + y_offset, 120, 80))

    def set_position(self, x, y):
        self.x = x
        self.y = y
        self.remake_rect()

class NodeContainer:
    def __init__(self):
        self.nodes: dict[UUID, Node] = {}

    def new_node(self, x: int, y: int):
        node = Node(x, y)
        uuid = new_uuid()
        self.nodes[uuid] = node
        return uuid

    def get_node(self, uuid: UUID) -> Node:
        return self.nodes[uuid]

    def draw(self, surface: pygame.Surface, node_color: pygame.Color, node_highlight_color: pygame.Color, text_color: pygame.Color):
        for uuid in self.nodes:
            node = self.nodes[uuid]
            node.draw(surface, node_color, node_highlight_color, text_color)
